import cv2
import pandas as pd
import csv
import numpy as np

import matplotlib.pyplot as plt
# from util.dataset_util import *
import torch
# from CustomImageDataset import CustomImageDataset
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda
import matplotlib.pyplot as plt
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        x = 128
        y = 64
        self.conv1 = nn.Conv2d(1, 8, 5, 1)
        x = (x-2)/2
        y = (y-2)/2
        self.conv2 = nn.Conv2d(8, 16, 3, 1)
        x = int((x-2)/2)
        y = int((y-2)/2)
        
        self.conv3 = nn.Conv2d(16, 32, 3, 1)
        x = int((x-2)/2)
        y = int((y-2)/2)
        self.conv4 = nn.Conv2d(32, 32, 3, 1)
        x = int((x-2)/2)
        y = int((y-2)/2)
        print(x, y)
        self.linear_relu_stack = nn.Sequential(
            # nn.Linear(32*32*3, 512),
            nn.Linear(x*y*32, 512),
            nn.ReLU(),
            # nn.Linear(512, 128),
            # nn.ReLU(),
            # nn.Linear(512, 128),
            # nn.ReLU(),
            # nn.Linear(256, 128),
            # nn.Sigmoid(),
            # nn.Linear(128, 32),
            # nn.Sigmoid(),
            nn.Linear(512, 3),
            # nn.ReLU(),
            # nn.Sigmoid()
        )
        self.flatten = nn.Flatten()

    # def forward(self, x):
    #     x = self.conv1(x)
    #     x = torch.nn.functional.relu(x)
    #     x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
    #     x = self.conv2(x)
    #     x = torch.nn.functional.relu(x)
    #     x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
    #     x = self.conv3(x)
    #     x = torch.nn.functional.relu(x)
    #     x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
    #     # print (x.shape)
    #     x = self.flatten(x)
    #     # print (x.shape)
    #     logits = self.linear_relu_stack(x)
    #     return logits
    def forward(self, x, show = False):
        if show: show_x = []
        x = self.conv1(x)
        if show: show_x.append(x.cpu())
        x = torch.nn.functional.relu(x)
        x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
        if show: show_x.append(x.cpu())
        x = self.conv2(x)
        if show: show_x.append(x.cpu())
        x = torch.nn.functional.relu(x)
        x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
        if show: show_x.append(x.cpu())
        x = self.conv3(x)
        x = torch.nn.functional.relu(x)
        if show: show_x.append(x.cpu())
        x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
        if show: show_x.append(x.cpu())
        x = self.conv4(x)
        x = torch.nn.functional.relu(x)
        if show: show_x.append(x.cpu())
        x = torch.nn.functional.max_pool2d(x, kernel_size = 2)
        if show: show_x.append(x.cpu())
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        if show:
            plt.figure(figsize=(15,2))
            plt.suptitle('hidden output', color='w', fontsize=32)
            for i in range(len(show_x)):
                plt.subplot(1,8,i+1)
                plt.imshow(show_x[i].cpu()[0][0].detach().numpy())
        return logits


def ellipse_detect(img):
    """
    :param image: 图片路径
    :return: None
    """
    # img = cv2.imread(image,cv2.IMREAD_COLOR)
    skinCrCbHist = np.zeros((256, 256), dtype= np.uint8 )
    cv2.ellipse(skinCrCbHist ,(113,155),(23,15),43,0, 360, (255,255,255),-1)
 
    YCRCB = cv2.cvtColor(img,cv2.COLOR_BGR2YCR_CB)
    (y,cr,cb)= cv2.split(YCRCB)
    skin = np.zeros(cr.shape, dtype=np.uint8)
    (x,y)= cr.shape
    for i in range(0,x):
        for j in range(0,y):
            CR= YCRCB[i,j,1]
            CB= YCRCB[i,j,2]
            if skinCrCbHist [CR,CB]>0:
                skin[i,j]= 255
    # cv2.namedWindow(image, cv2.WINDOW_NORMAL)
    # cv2.imshow(image, img)
    dst = cv2.bitwise_and(img,img,mask= skin)
    return dst

def test_mp4(src, model):
    video = cv2.VideoCapture(src)
    # plt.figure()
    cnt = 0
    while True:
        ret, img = video.read() # img 就是一帧图片
        
        if not ret: break # 当获取完最后一帧就结束
        # img = img(cv2.Range(100, 1280-100), cv2.Range(100,720-100))
        img = cv2.resize(img, (128,64))
        # if train_test == 'test': img = img[100: 720-100,100:1280-100] 
        img = ellipse_detect(img)

        # img = cv2.resize(img, (128,64))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # img = img.astype(np.float32)
        img = cv2.Canny(img, 220, 255)
        img = img.astype(np.float32)
        img = torch.Tensor(img)
        img = img.unsqueeze(0).unsqueeze(0).cuda().float()
        print(torch.nn.functional.softmax(model(img), dim = 1).argmax().item())
        # _, canny = cv2.threshold(canny, 220, 255, cv2.THRESH_BINARY) 
        # plt.imshow(canny)
        # plt.show()
        # canny = canny.astype(np.float32)
        # img/=256
        # cv2.imshow('img', img)
        # cv2.imshow('canny', canny)
        # cv2.waitKey(10)
        
    cv2.destroyAllWindows()

model = torch.load('predict_model.pkl').cuda()

test_mp4('raw_data/test/one.mp4', model)